import socket
import json
from shared.crypto import CryptoUtils
import os
import cmd
import getpass
import threading

class FTPClient(cmd.Cmd):
    """
    FTP客户端类
    继承自cmd.Cmd，提供命令行界面：
    - 支持与服务器的加密通信
    - 实现文件传输功能
    - 支持聊天功能
    - 提供目录操作命令
    """

    intro = '欢迎来到FTP客户端。输入help或者？查看命令.\n'
    prompt = 'ftp> '

    def __init__(self):
        """
        初始化FTP客户端
        - 初始化基类
        - 设置连接状态
        - 初始化消息处理相关的变量
        """
        super().__init__()
        self.socket = None
        self.symmetric_key = None
        self.connected = False
        self.authenticated = False
        self.message_thread = None
        self.response_queue = []  # 用于存储命令响应
        self.message_lock = threading.Lock()  # 用于同步消息处理

    def connect(self, host='localhost', port=21000):
        """
        连接到FTP服务器
        - 建立TCP连接
        - 完成密钥交换
        Args:
            host (str): 服务器地址
            port (int): 服务器端口
        Returns:
            bool: 是否连接成功
        """
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect((host, port))
            
            # 1.接收服务器的RSA公钥
            public_key_pem = self.socket.recv(2048)
            public_key = CryptoUtils.deserialize_public_key(public_key_pem)

            # 2.生成对称密钥，使用服务器公钥加密这个对称密钥
            self.symmetric_key = CryptoUtils.generate_symmetric_key()
            encrypted_symmetric_key = CryptoUtils.encrypt_with_public_key(
                public_key, self.symmetric_key)
            # 3.发送加密后的对称密钥到服务器（服务端使用私钥解密之后就能拿到相同的对称密钥，之后就能用这个对称密钥进行安全传输）
            self.socket.send(encrypted_symmetric_key)

            self.connected = True
            return True
        except Exception as e:
            print(f"Connection failed: {str(e)}")
            return False

    def wait_for_response(self, timeout=5):
        """
        等待服务器响应
        实现超时机制防止无限等待
        Args:
            timeout (int): 超时时间（秒）
        Returns:
            dict: 服务器响应数据，超时返回None
        """
        import time
        start_time = time.time()
        while time.time() - start_time < timeout:
            with self.message_lock:
                if self.response_queue:
                    return self.response_queue.pop(0)
            time.sleep(0.1)
        return None

    def send_command(self, command, args=None):
        """
        发送命令到服务器
        - 检查认证状态
        - 加密命令数据
        - 发送并等待响应
        Args:
            command (str): 命令名称
            args (list, optional): 命令参数
        Returns:
            dict: 服务器响应数据
        """
        if not self.authenticated:
            print("Please authenticate first")
            return None

        command_data = {
            'command': command,
            'args': args or []
        }
        encrypted_data = CryptoUtils.encrypt_message(
            self.symmetric_key,
            json.dumps(command_data)
        )
        try:
            self.socket.send(encrypted_data)
            
            # 等待响应
            response = self.wait_for_response()
            if not response:
                if command == 'chat':
                    # 对于聊天命令，超时不显示错误
                    return {'success': True}
                return {'success': False, 'message': '服务器未响应'}
            return response
        except Exception as e:
            return {'success': False, 'message': f'发送失败: {str(e)}'}

    def start_message_listener(self):
        """
        启动消息监听线程
        - 持续监听服务器消息
        - 处理聊天消息和命令响应
        - 使用线程锁保证线程安全
        """
        def message_listener():
            while self.authenticated:
                try:
                    encrypted_data = self.socket.recv(4096)
                    if not encrypted_data:
                        break
                    
                    decrypted_data = CryptoUtils.decrypt_message(
                        self.symmetric_key, encrypted_data)
                    data = json.loads(decrypted_data)
                    
                    with self.message_lock:
                        if isinstance(data, dict):
                            if data.get('type') == 'chat':
                                # 处理聊天消息
                                print(f"\n[来自 {data['from']} 的消息]: {data['message']}")
                                print(self.prompt, end='', flush=True)
                            else:
                                # 存储命令响应
                                self.response_queue.append(data)
                except Exception as e:
                    if self.authenticated:
                        print(f"\n接收消息时出错: {str(e)}")
                        print(self.prompt, end='', flush=True)
                    break

        import threading
        if self.message_thread is None:
            self.message_thread = threading.Thread(target=message_listener)
            self.message_thread.daemon = True
            self.message_thread.start()

    def authenticate(self):
        """
        进行用户认证
        - 获取用户凭据
        - 加密发送到服务器
        - 处理认证响应
        Returns:
            bool: 认证是否成功
        """
        if not self.connected:
            print("Not connected to server")
            return False

        username = input("Username: ")
        password = getpass.getpass("Password: ")

        credentials = {
            'username': username,
            'password': password
        }
        encrypted_credentials = CryptoUtils.encrypt_message(
            self.symmetric_key,
            json.dumps(credentials)
        )
        self.socket.send(encrypted_credentials)

        encrypted_response = self.socket.recv(1024)
        response = json.loads(CryptoUtils.decrypt_message(
            self.symmetric_key,
            encrypted_response
        ))

        if response['success']:
            print("Authentication successful")
            self.authenticated = True
            self.start_message_listener()  # 开始监听消息
            return True
        else:
            print("Authentication failed")
            return False

    def do_connect(self, arg):
        """
        处理connect命令
        连接到指定的服务器并进行认证
        Usage: connect [host] [port]
        """
        args = arg.split()
        host = args[0] if args else 'localhost'
        port = int(args[1]) if len(args) > 1 else 21000
        if self.connect(host, port):
            if self.authenticate():
                print(f"Connected to {host}:{port}")
            else:
                self.socket.close()
                self.connected = False
        
    def do_list(self, arg):
        """
        处理list命令
        列出当前目录下的文件和目录
        Usage: list
        """
        response = self.send_command('list')
        if response and response['success']:
            for item in response['data']:
                print(item)
        else:
            print("Error:", response.get('message', 'Unknown error'))

    def do_cd(self, arg):
        """
        处理cd命令
        切换当前工作目录
        Usage: cd [path]
        """
        response = self.send_command('cd', [arg] if arg else [])
        if not response['success']:
            print("Error:", response.get('message', 'Unknown error'))

    def do_pwd(self, arg):
        """
        处理pwd命令
        显示当前工作目录
        Usage: pwd
        """
        response = self.send_command('pwd')
        if response and response['success']:
            print(response['data'])
        else:
            print("Error:", response.get('message', 'Unknown error'))

    def do_mkdir(self, arg):
        """
        处理mkdir命令
        创建新目录
        Usage: mkdir <dirname>
        """
        if not arg:
            print("Directory name required")
            return
        response = self.send_command('mkdir', [arg])
        if not response['success']:
            print("Error:", response.get('message', 'Unknown error'))

    def do_get(self, arg):
        """
        处理get命令
        从服务器下载文件
        Usage: get <filename>
        """
        if not arg:
            print("Filename required")
            return
        response = self.send_command('get', [arg])
        if response and response['success']:
            try:
                data = bytes.fromhex(response['data'])
                with open(arg, 'wb') as f:
                    f.write(data)
                print(f"File '{arg}' downloaded successfully")
            except Exception as e:
                print(f"Error saving file: {str(e)}")
        else:
            print("Error:", response.get('message', 'Unknown error'))

    def do_put(self, arg):
        """
        处理put命令
        将文件上传到服务器
        Usage: put <filename>
        """
        if not arg:
            print("Filename required")
            return
        try:
            with open(arg, 'rb') as f:
                data = f.read()
            response = self.send_command('put', [os.path.basename(arg), data.hex()])
            if response and response['success']:
                print(f"File '{arg}' uploaded successfully")
            else:
                print("Error:", response.get('message', 'Unknown error'))
        except Exception as e:
            print(f"Error reading file: {str(e)}")

    def do_attrib(self, arg):
        """
        处理attrib命令
        查看文件的属性信息
        Usage: attrib <filename>
        """
        if not arg:
            print("Filename required")
            return
        response = self.send_command('attrib', [arg])
        if response and response['success']:
            attrs = response['data']
            print(f"File: {arg}")
            print(f"Size: {attrs['size']} bytes")
            print(f"Created: {attrs['created']}")
            print(f"Modified: {attrs['modified']}")
            print(f"Mode: {attrs['mode']}")
        else:
            print("Error:", response.get('message', 'Unknown error'))

    def do_rename(self, arg):
        """
        处理rename命令
        重命名文件或目录
        Usage: rename <old_name> <new_name>
        """
        args = arg.split()
        if len(args) != 2:
            print("Old and new names required")
            return
        response = self.send_command('rename', args)
        if not response['success']:
            print("Error:", response.get('message', 'Unknown error'))

    def do_create(self, arg):
        """
        处理create命令
        创建空文件
        Usage: create <filename>
        """
        if not arg:
            print("Filename required")
            return
        response = self.send_command('create', [arg])
        if not response['success']:
            print("Error:", response.get('message', 'Unknown error'))

    def do_add(self, arg):
        """
        处理add命令
        添加新用户（需要管理员权限）
        Usage: add <username> <password>
        """
        if not self.connected:
            print("Please connect to server first")
            return
        args = arg.split()
        if len(args) != 2:
            print("Username and password required")
            return
        response = self.send_command('add', args)
        if response is None:
            return  # send_command已经打印了认证错误消息
        if response.get('success', False):
            print("User added successfully")
        else:
            print("Error:", response.get('message', 'Unknown error'))

    def do_chat(self, arg):
        """
        处理chat命令
        发送聊天消息到指定用户或广播给所有用户
        Usage: chat <recipient> <message>
        - recipient可以是用户名或"all"（广播）
        """
        if not arg:
            print("Usage: chat <recipient> <message>")
            return
        
        try:
            parts = arg.split(maxsplit=1)
            if len(parts) < 2:
                print("请同时提供接收者和消息内容")
                return
                
            recipient, message = parts
            response = self.send_command('chat', [recipient, message])
            if response:
                if response['success']:
                    print("消息已发送")
                else:
                    print("发送失败:", response.get('message', '未知错误'))
        except Exception as e:
            print(f"发送消息时出错: {str(e)}")

    def do_online(self, arg):
        """
        处理online命令
        获取当前在线用户列表
        Usage: online
        """
        try:
            response = self.send_command('online')
            if response:
                if response['success'] and 'data' in response:
                    users = response['data']
                    if users:
                        print("\n在线用户:")
                        for user in users:
                            print(f"- {user}")
                    else:
                        print("当前没有其他用户在线")
                else:
                    print("获取在线用户列表失败:", response.get('message', '未知错误'))
        except Exception as e:
            print(f"获取在线用户列表时出错: {str(e)}")

    def do_quit(self, arg):
        """
        处理quit命令
        关闭客户端连接并退出程序
        Usage: quit
        """
        if self.socket:
            self.authenticated = False
            if self.message_thread:
                self.message_thread.join(timeout=1)
            self.socket.close()
        return True

if __name__ == '__main__':
    FTPClient().cmdloop()
